import matplotlib as mpl
mpl.use('Agg')
from  MaxEnt import MaxEnt
import argparse
import os,glob
from collections import namedtuple
from multiprocessing import Pool
import matplotlib.pyplot as plt
from records import recordtype 
from OrganizeData import OrganizeMaxEntData
from MaxEntParams import mxParams
import help_func
from plotMaxEnt import PlotSaveMaxEnt
import mkl
import tables as tb
import math
import sys
import numpy as np
from ACSmooth import ACLap


class ScaledLocator(mpl.ticker.MaxNLocator):
    """
    Locates regular intervals along an axis scaled by *dx* and shifted by
    *x0*. For example, this would locate minutes on an axis plotted in seconds
    if dx=60.  This differs from MultipleLocator in that an approriate interval
    of dx units will be chosen similar to the default MaxNLocator.
    """
    def __init__(self, dx=1.0, x0=0.0):
        self.dx = dx
        self.x0 = x0
        mpl.ticker.MaxNLocator.__init__(self, nbins=9, steps=[1, 2, 5, 10])

    def rescale(self, x):
        return x / self.dx + self.x0
    def inv_rescale(self, x):
        return  (x - self.x0) * self.dx

    def __call__(self): 
        vmin, vmax = self.axis.get_view_interval()
        vmin, vmax = self.rescale(vmin), self.rescale(vmax)
        vmin, vmax = mpl.transforms.nonsingular(vmin, vmax, expander = 0.05)
        locs = self.bin_boundaries(vmin, vmax)
        locs = self.inv_rescale(locs)
        prune = self._prune
        if prune=='lower':
            locs = locs[1:]
        elif prune=='upper':
            locs = locs[:-1]
        elif prune=='both':
            locs = locs[1:-1]
        return self.raise_if_exceeds(locs)

class ScaledFormatter(mpl.ticker.OldScalarFormatter):
    """Formats tick labels scaled by *dx* and shifted by *x0*."""
    def __init__(self, dx=1.0, x0=0.0, **kwargs):
        self.dx, self.x0 = dx, x0

    def rescale(self, x):
        return x / self.dx + self.x0

    def __call__(self, x, pos=None):
        xmin, xmax = self.axis.get_view_interval()
        xmin, xmax = self.rescale(xmin), self.rescale(xmax)
        d = abs(xmax - xmin)
        x = self.rescale(x)
        s = self.pprint_val(x, d)
        return s

def listDirs(in_dir):  
    dirs=[d for d in glob.glob(in_dir+'/???') if os.path.isdir(d)]
    print dirs
    return dirs
def maxEntSingleDir(mxData, out_dir,dump_dir):
    mxData.dir = out_dir
    mxData.dump_dir = out_dir+"/dump_"+dump_dir
    if not os.path.exists(mxData.dump_dir):
        os.makedirs(mxData.dump_dir)    
    od = OrganizeMaxEntData()
    #od.check_eigs(mxData, "VER2")
    simData = od.readData(mxData, "VER0")
    maxEnt = MaxEnt()
    runResults, mxData = maxEnt.run(simData, mxData)
    psMaxEnt = PlotSaveMaxEnt()
    acLap=ACLap(simData, mxData)
    allRes=acLap.run()
    acLap.printAll(allRes)
    return psMaxEnt.plotSaveData(runResults, mxData)
def dump(anData, out_dir):
    outFile = out_dir + "/MaxEnt.h5"
    h5file = tb.openFile(outFile, "w")
    for i, data in enumerate(anData):
        cgroup = h5file.createGroup("/", "sim_" + str(i))
        h5file.createArray(cgroup, "maxA", data["anData"]["A_maxProb"])
        h5file.createArray(cgroup, "maxCurvatureA", data["anData"]["A_maxCurve"])
        h5file.createArray(cgroup, "G", data["anData"]["G"]) 
        h5file.createArray(cgroup, "omega", data["mxData"].omegaP)
    h5file.close()
def plotAll(anData, out_dir,mxData):
    fig = plt.figure()
    colormap = plt.cm.gist_ncar
    ax1 = fig.add_subplot(211)
    ax2 = fig.add_subplot(212)
    ax1.set_color_cycle([colormap(i) for i in np.linspace(0, 0.9, 10)])
    ax2.set_color_cycle([colormap(i) for i in np.linspace(0, 0.9, 10)])
    deltadis=mxData.deltapos
    betazero=120.	
    for data in anData:
        gc=data["mxData"].gc
        muc=data["mxData"].muc
        #deltag=math.fabs((1-float(data["anData"]["G"])/gc))
    	deltag=math.fabs((1-float(data["anData"]["MU"])/muc))
    #ax.plot(self.mxData.omegaP, np.multiply(self.mxData.omegaP, anp))
    	ax1.plot(data["mxData"].omegaP, data["anData"]["A_maxProb"], label=r'$\delta g={0:.2f}$'.format(100*deltag))
    	ax2.plot(data["mxData"].omegaP, data["anData"]["A_maxCurve"], label=r'$\delta g={0:.2f}$'.format(100*deltag))
    ax1.set_xlim(0.,3.)
    ax2.set_xlim(0.,3.)   
    handles, labels = ax1.get_legend_handles_labels()
    # reverse the order
    ax1.legend(handles[::-1], labels[::-1],prop={'size':6})
    leg = plt.gca().get_legend()
    handles, labels = ax2.get_legend_handles_labels()       # reverse the order
    ax2.legend(handles[::-1], labels[::-1],prop={'size':6})
    fig.savefig(out_dir + "/allA.pdf")
    #Scaled
    fig = plt.figure()
    ax1 = fig.add_subplot(211)
    ax2 = fig.add_subplot(212)
    ax1.set_color_cycle([colormap(i) for i in np.linspace(0, 0.9, 10)])
    ax2.set_color_cycle([colormap(i) for i in np.linspace(0, 0.9, 10)])
    for data in anData:
        #ax.plot(self.mxData.omegaP, np.multiply(self.mxData.omegaP, anp))
        gc=data["mxData"].gc
        muc=data["mxData"].muc
#	    deltag=math.fabs((1-float(data["anData"]["G"])/gc))
    	deltag=math.fabs((1-float(data["anData"]["MU"])/muc))
    	deltascale=(deltag**data["mxData"].nu)*deltadis	
    	ax1.plot(data["mxData"].omegaP*data["anData"]["beta"]/betazero, data["anData"]["A_maxProb"], label=r'$\delta g={0:.2f}$'.format(100*deltag))
    	ax2.plot(data["mxData"].omegaP*data["anData"]["beta"]/betazero, data["anData"]["A_maxCurve"], label=r'$\delta g={0:.2f}$'.format(100*deltag))
    ax1.set_xlim(0.,3.)
    ax2.set_xlim(0.,3.)
    handles, labels = ax1.get_legend_handles_labels()
    # reverse the order
    ax1.legend(handles[::-1], labels[::-1],prop={'size':6})
    leg = plt.gca().get_legend()
    handles, labels = ax2.get_legend_handles_labels()       # reverse the order
    ax2.legend(handles[::-1], labels[::-1],prop={'size':6})
    fig.savefig(out_dir + "/allAScalled.pdf")
#Scaled
    fig = plt.figure()
    ax1 = fig.add_subplot(111)
    ax1.set_color_cycle([colormap(i) for i in np.linspace(0, 0.9, 10)])
    for data in anData:
        gc=data["mxData"].gc	
        #ax.plot(self.mxData.omegaP, np.multiply(self.mxData.omegaP, anp))
        muc=data["mxData"].muc
#	    deltag=math.fabs((1-float(data["anData"]["G"])/gc))
    	deltag=math.fabs((1-float(data["anData"]["MU"])/muc))
    	deltascale=(deltag**data["mxData"].nu)*deltadis		
    	ax1.plot(data["mxData"].omegaP*data["anData"]["beta"]/betazero, data["anData"]["A_maxCurve"], label=r'$\delta g={0:.2f}$'.format(100*deltag))
    ax1.set_xlim(0., 3.)
    handles, labels = ax1.get_legend_handles_labels()
    # reverse the order
    ax1.legend(handles[::-1], labels[::-1],prop={'size':10})
    leg = plt.gca().get_legend()
    fig.savefig(out_dir + "/allAScalled2.pdf")
if __name__ == "__main__":
    mkl.set_num_threads(4)
    # arguments
    parser = argparse.ArgumentParser(description='MaxEnt')
    parser.add_argument('-dir', action='store')
    parser.add_argument('-sims', nargs='*', type=int, action='store')
    parser.add_argument('-dump', action='store',default="")
    args = parser.parse_args()
    #sys.stdout = open(args.dir+'/MaxEnt.txt', 'w')
    # find dirs
    sim_dirs = listDirs(args.dir)
    mxData = help_func.MaxEntAlgData(**mxParams)        
    #sim_dirs=["/srv01/technion/snirgaz/Desktop"]
    anData = map(lambda x:maxEntSingleDir(mxData, x,args.dump), sim_dirs)
    dump_dir=args.dir+"/dump_"+args.dump
    if not os.path.exists(dump_dir):
    	os.makedirs(dump_dir) 
    plotAll(anData, dump_dir,mxData)
    dump(anData, dump_dir)
    # Organize Data
    
    
    
    #
#    
#    omegaMax = 15
#    numOfOmega = 1000
#    alphaMin = 0
#    alphaMax = 500000
#    errorLevel= 3
#    MaxEntAlgData = recordtype('MaxEntCalcData', \
#                                ['dir', \
#                                 'omegaMax', \
#                                 'numOfOmega', \
#                                 'alphaMin', \
#                                 'alphaMax',\
#                                 'errorLevel'])
#    for dir in sim_dirs:
#        # Organize data
#        
#        
#        # run Max Ent
#        
#    MaxEntSims = [MaxEntAlg(MaxEntAlgData(dir=d, omegaMax=omegaMax, numOfOmega=numOfOmega, obs=obs, alphaMin=alphaMin, alphaMax=alphaMax,errorLevel=errorLevel)) for d in simDir]
#    
#    for MaxEntSim in MaxEntSims:
#        MaxEntSim.run()
#    fAll=plt.figure(1)
#    fAll.savefig(args.dir+"/sw.pdf",dpi=100)
    
    
    
        
